#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/image.hpp>
#include <astra/capi/astra.h>
#include <astra/capi/streams/infrared_capi.h>
#include <astra/capi/streams/color_capi.h>
#include <opencv2/opencv.hpp>
#include <chrono>
#include <sys/stat.h>
#include <thread>
#include <std_msgs/msg/float32_multi_array.hpp>
#include <std_msgs/msg/float64_multi_array.hpp>
#include <mutex>
#include <iomanip>
#include <algorithm>
#include <std_msgs/msg/string.hpp>

using namespace std::chrono_literals;
namespace
{
    constexpr int FIXED_WIDTH = 640;
    constexpr int FIXED_HEIGHT = 480;
    constexpr int FIXED_LENGTH = 307200;
    constexpr int FIXED_LENGTH_3 = 921600;
    alignas(64) uint8_t image_buffer[FIXED_LENGTH_3] = {0};
    uint32_t ir_green_indices[FIXED_LENGTH];
    uint32_t g_green_count = 0;
    uint8_t *ir_data;
    uint32_t ir_length;
    bool fire_detected_ir;
    bool fire_detected_yolo;
}

class AstraNode : public rclcpp::Node
{
public:
    AstraNode() : Node("astra_node"), running_(true)
    {
        // 初始化彩色图像发布器
        publisher_ = this->create_publisher<sensor_msgs::msg::Image>("/camera", 10);

        video_path_publisher_ = this->create_publisher<std_msgs::msg::String>("/video_path", 10);

        detection_sub_ = this->create_subscription<std_msgs::msg::Float64MultiArray>(
            "/fire_yolo_detection", 10,
            [this](const std_msgs::msg::Float64MultiArray::SharedPtr msg)
            {
                std::lock_guard<std::mutex> lock(detection_mutex_);
                if (msg->data.size() != 7)
                {
                    RCLCPP_ERROR(this->get_logger(), "Invalid detection data size: expected 7, got %zu", msg->data.size());
                    return;
                }

                DetectionInfo info;
                info.x1 = static_cast<float>(msg->data[0]);
                info.y1 = static_cast<float>(msg->data[1]);
                info.x2 = static_cast<float>(msg->data[2]);
                info.y2 = static_cast<float>(msg->data[3]);
                info.confidence = static_cast<float>(msg->data[4]);

                // 直接存储为纳秒时间戳
                int64_t sec = static_cast<int64_t>(msg->data[5]);
                int64_t nsec = static_cast<int64_t>(msg->data[6]);
                info.stamp_ns = sec * 1000000000LL + nsec;

                latest_detections_.push_back(info);
            });

        // 初始化Astra设备
        if (astra_initialize() != ASTRA_STATUS_SUCCESS)
        {
            RCLCPP_ERROR(this->get_logger(), "Failed to initialize Astra");
            rclcpp::shutdown();
        }
        if (astra_streamset_open("device/default", &sensor_) != ASTRA_STATUS_SUCCESS)
        {
            RCLCPP_ERROR(this->get_logger(), "Failed to open sensor");
            astra_terminate();
            rclcpp::shutdown();
        }
        // 创建数据流读取器
        if (astra_reader_create(sensor_, &reader_) != ASTRA_STATUS_SUCCESS)
        {
            RCLCPP_ERROR(this->get_logger(), "Failed to create reader");
            astra_streamset_close(&sensor_);
            astra_terminate();
            rclcpp::shutdown();
        }
        // 配置并启动流
        configure_streams();
        // 禁用激光器
        astra_status_t status;
        if ((status = astra_ldp_enable(sensor_, false)) != ASTRA_STATUS_SUCCESS ||
            (status = astra_laser_enable(sensor_, false)) != ASTRA_STATUS_SUCCESS)
        {
            std::cerr << "Error disabling laser: " << status << std::endl;
        }

        // 启动数据处理线程
        processing_thread_ = std::thread(&AstraNode::processing_loop, this);
    }

    ~AstraNode()
    {
        running_ = false;
        if (processing_thread_.joinable())
        {
            processing_thread_.join();
        }
        if (video_path_publish_10times_thread_.joinable())
        {
            video_path_publish_10times_thread_.join();
        }
        astra_reader_destroy(&reader_);
        astra_streamset_close(&sensor_);
        astra_terminate();
    }

private:
    void processing_loop()
    {
        auto last_process_time = std::chrono::steady_clock::now();
        astra_reader_frame_t frame;

        while (rclcpp::ok() && running_)
        {
            auto now = std::chrono::steady_clock::now();
            if (std::chrono::duration_cast<std::chrono::milliseconds>(now - last_process_time).count() >= 96)
            {
                astra_update();
                if (astra_reader_open_frame(reader_, 0, &frame) == ASTRA_STATUS_SUCCESS)
                {
                    {
                        std::lock_guard<std::mutex> lock(buffer_mutex_);
                        memset(image_buffer, 0, sizeof(image_buffer));
                        // 处理彩色帧
                        astra_colorframe_t color_frame;
                        if (astra_frame_get_colorframe(frame, &color_frame) == ASTRA_STATUS_SUCCESS)
                        {
                            process_color(color_frame);
                        }
                        // 处理红外帧
                        // astra_infraredframe_t ir_frame;
                        // if (astra_frame_get_infraredframe(frame, &ir_frame) == ASTRA_STATUS_SUCCESS) {
                        //     process_infrared(ir_frame);
                        // }
                    }

                    cv::Mat image_visual_(FIXED_HEIGHT, FIXED_WIDTH, CV_8UC3, image_buffer);
                    {
                        std::lock_guard<std::mutex> lock(detection_mutex_);
                        // 获取当前时间戳（转换为纳秒）
                        int64_t now_ns = this->get_clock()->now().nanoseconds();

                        // 移除过期检测结果（0.2秒前）
                        latest_detections_.erase(
                            std::remove_if(latest_detections_.begin(), latest_detections_.end(),
                                           [now_ns](const DetectionInfo &info)
                                           {
                                               return (now_ns - info.stamp_ns) > 200000000LL; // 0.2秒=200,000,000纳秒
                                           }),
                            latest_detections_.end());
                        fire_detected_yolo = false;
                        // 绘制检测框
                        for (const auto &detection : latest_detections_)
                        {
                            fire_detected_yolo = true;
                            // 绘制矩形
                            cv::rectangle(image_visual_,
                                          cv::Point((int)detection.x1, (int)(detection.y1 * FIXED_HEIGHT / FIXED_WIDTH)),
                                          cv::Point((int)detection.x2, (int)(detection.y2 * FIXED_HEIGHT / FIXED_WIDTH)),
                                          cv::Scalar(0, 255, 0), 2);
                            // 绘制置信度
                            std::stringstream ss;
                            ss << std::fixed << std::setprecision(2) << detection.confidence;
                            cv::putText(image_visual_, ss.str(),
                                        cv::Point((int)detection.x1 + 2, (int)detection.y1 + 20),
                                        cv::FONT_HERSHEY_SIMPLEX, 0.6,
                                        cv::Scalar(0, 255, 0), 1);
                        }
                    }

                    // 录像逻辑
                    // handle_recording();
                    handle_save_image();
                    // 显示图像
                    // cv::imshow("Infrared Visual", image_visual_);
                    // cv::waitKey(1);
                }
                astra_reader_close_frame(&frame);
                last_process_time = now;
            }
            else
            {
                std::this_thread::sleep_for(20ms);
            }
        }
    }

    void configure_streams()
    {
        // astra_reader_get_infraredstream(reader_, &infrared_stream_);
        // astra_imagestream_mode_t ir_mode;
        // ir_mode.width = 640;
        // ir_mode.height = 480;
        // ir_mode.pixelFormat = ASTRA_PIXEL_FORMAT_GRAY16;
        // ir_mode.fps = 30;
        // astra_imagestream_set_mode(infrared_stream_, &ir_mode);
        // astra_stream_start(infrared_stream_);

        astra_reader_get_colorstream(reader_, &color_stream_);
        astra_imagestream_mode_t color_mode;
        color_mode.width = 640;
        color_mode.height = 480;
        color_mode.pixelFormat = ASTRA_PIXEL_FORMAT_RGB888;
        color_mode.fps = 30;
        astra_imagestream_set_mode(color_stream_, &color_mode);
        astra_stream_start(color_stream_);
    }

    void process_infrared(astra_infraredframe_t frame)
    {
        astra_status_t status = astra_infraredframe_get_data_ptr(frame, &ir_data, &ir_length);
        if (status != ASTRA_STATUS_SUCCESS || !ir_data)
            return;

        g_green_count = 0;
        bool has_high = false;
        for (int i = 0, j = FIXED_HEIGHT - 1; i < FIXED_HEIGHT; i++, j--)
        {
            for (int k = 0; k < FIXED_WIDTH; k++)
            {
                const uint8_t red = (ir_data[(j * FIXED_WIDTH + k) * 2] >> 2) + (ir_data[(j * FIXED_WIDTH + k) * 2 + 1] << 6);
                image_buffer[FIXED_LENGTH_3 + (i * FIXED_WIDTH + k) * 3 + 2] = red;
                has_high |= (red > 225);
                if (120 <= red && red <= 160)
                {
                    ir_green_indices[g_green_count++] = FIXED_LENGTH_3 + (i * FIXED_WIDTH + k) * 3;
                }
            }
        }
        fire_detected_ir = has_high;
        if (has_high)
        {
            for (int i = 0; i < g_green_count; i++)
            {
                image_buffer[ir_green_indices[i]] = 255;
                image_buffer[ir_green_indices[i] + 1] = 255;
                image_buffer[ir_green_indices[i] + 2] = 255;
            }
        }
    }

    void process_color(astra_colorframe_t frame)
    {

        uint8_t *color_data;
        uint32_t color_length;
        astra_colorframe_get_data_ptr(frame, &color_data, &color_length);
        for (int i = 0, j = FIXED_HEIGHT - 1; i < FIXED_HEIGHT; i++, j--)
        {
            for (int k = 0; k < FIXED_WIDTH; k++)
            {
                image_buffer[(i * FIXED_WIDTH + k) * 3] = color_data[(j * FIXED_WIDTH + k) * 3 + 2];
                image_buffer[(i * FIXED_WIDTH + k) * 3 + 1] = color_data[(j * FIXED_WIDTH + k) * 3 + 1];
                image_buffer[(i * FIXED_WIDTH + k) * 3 + 2] = color_data[(j * FIXED_WIDTH + k) * 3];
            }
        }
        cv::Mat color_img(FIXED_HEIGHT, FIXED_WIDTH, CV_8UC3, image_buffer);

        auto now = std::chrono::steady_clock::now();
        if (std::chrono::duration_cast<std::chrono::milliseconds>(now - last_color_post_time).count() >= 96)
        {
            auto msg = std::make_shared<sensor_msgs::msg::Image>();
            msg->header.stamp = this->now();
            msg->header.frame_id = "camera";
            msg->height = color_img.rows;
            msg->width = color_img.cols;
            msg->encoding = "bgr8";
            msg->step = color_img.step;
            msg->data.assign(color_img.data, color_img.data + color_img.total() * color_img.elemSize());
            publisher_->publish(*msg);
            last_color_post_time = now;
        }
    }

    void handle_recording()
    {
        auto now = std::chrono::steady_clock::now();
        if (fire_detected_ir || fire_detected_yolo)
        {
            last_fire_time_ = now;
            if (!is_recording_)
                start_recording();
        }
        else if (is_recording_ && std::chrono::duration_cast<std::chrono::seconds>(now - last_fire_time_).count() > 5)
        {
            stop_recording();
        }
        if (is_recording_)
        {
            cv::Mat image_visual_(FIXED_HEIGHT, FIXED_WIDTH, CV_8UC3, image_buffer);
            video_writer_.write(image_visual_);
        }
    }

    void start_recording()
    {
        std::string path = std::string(getenv("HOME")) + "/fire_detection";
        mkdir(path.c_str(), 0777);
        auto t = std::time(nullptr);
        std::stringstream ss;
        ss << std::put_time(std::localtime(&t), "%Y%m%d_%H%M%S");
        current_video_path_ = path + "/" + ss.str() + ".avi"; // 保存完整路径
        video_writer_.open(current_video_path_,
                           cv::VideoWriter::fourcc('M', 'J', 'P', 'G'), 10,
                           cv::Size(FIXED_WIDTH, FIXED_HEIGHT));
        is_recording_ = true;
    }

    void stop_recording()
    {
        if (video_writer_.isOpened())
        {
            video_writer_.release();
        }
        if (!current_video_path_.empty())
        {
            // 等待已有线程完成
            if (video_path_publish_10times_thread_.joinable())
            {
                video_path_publish_10times_thread_.join();
            }
            // 启动新线程
            video_path_publish_10times_thread_ = std::thread(&AstraNode::video_path_publish_10times, this);
        }else{
            video_path_publish_10times_thread_ = std::thread(&AstraNode::video_path_publish_10times, this);

        }
        is_recording_ = false;
    }

    void video_path_publish_10times()
    {
        auto msg = std_msgs::msg::String();
        msg.data = current_video_path_;
        for (int i = 0; i < 10; i++)
        {
            video_path_publisher_->publish(msg);
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
    }

    std::mutex buffer_mutex_;
    std::mutex saving_mutex_;
    bool is_saving_images_ = false;
    std::string image_folder_path_;



    void handle_save_image()
    {
        if (fire_detected_ir || fire_detected_yolo)
        {
            std::lock_guard<std::mutex> lock(saving_mutex_);
            if (!is_saving_images_)
            {
                start_save_image();
                is_saving_images_ = true;
                std::thread([this]() {
                    for (int i = 0; i < 5; ++i)
                    {
                        save_current_image();
                        if (i < 4) std::this_thread::sleep_for(400ms);
                    }
                    publish_image_folder();
                    std::this_thread::sleep_for(5000ms);
                    {
                        std::lock_guard<std::mutex> lock(saving_mutex_);
                        is_saving_images_ = false;
                    }
                }).detach();
            }
        }
    }

    void start_save_image()
    {
        // 创建父目录
        std::string parent_dir = std::string(getenv("HOME")) + "/fire_detection";
        struct stat st;
        if (stat(parent_dir.c_str(), &st) != 0) {
            mkdir(parent_dir.c_str(), 0777);
        }
        
        // 创建时间戳子目录
        auto t = std::time(nullptr);
        std::stringstream ss;
        ss << std::put_time(std::localtime(&t), "%Y%m%d_%H%M%S");
        image_folder_path_ = parent_dir + "/" + ss.str();
        mkdir(image_folder_path_.c_str(), 0777);
    }

    void save_current_image()
    {

        auto now = std::chrono::system_clock::now();
        auto now_time_t = std::chrono::system_clock::to_time_t(now);
        auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()) % 1000;

        std::stringstream filename;
        filename << std::put_time(std::localtime(&now_time_t), "%Y%m%d_%H%M%S")
                 << "_" << std::setfill('0') << std::setw(3) << ms.count() << ".jpg";
        
        std::string file_path = image_folder_path_ + "/" + filename.str();

            cv::Mat color_img;
    {
        std::lock_guard<std::mutex> lock(buffer_mutex_);
        color_img = cv::Mat(FIXED_HEIGHT, FIXED_WIDTH, CV_8UC3, image_buffer); // 直接引用缓冲区
        cv::imwrite(file_path, color_img); // 在锁内执行保存
    }
    }


    void publish_image_folder()
    {
        auto msg = std_msgs::msg::String();
        msg.data = image_folder_path_;
        for (int i = 0; i < 6; ++i)
        {
            video_path_publisher_->publish(msg);
            std::this_thread::sleep_for(333ms);
        }
    }

    // 成员变量
    rclcpp::Publisher<sensor_msgs::msg::Image>::SharedPtr publisher_;
    astra_streamsetconnection_t sensor_;
    astra_reader_t reader_;
    astra_infraredstream_t infrared_stream_;
    astra_colorstream_t color_stream_;
    std::thread processing_thread_;
    std::thread video_path_publish_10times_thread_;
    bool running_;
    std::chrono::steady_clock::time_point last_color_post_time;
    bool is_recording_ = false;
    cv::VideoWriter video_writer_;
    std::chrono::steady_clock::time_point last_fire_time_;

    rclcpp::Publisher<std_msgs::msg::String>::SharedPtr video_path_publisher_;
    std::string current_video_path_;

    rclcpp::Subscription<std_msgs::msg::Float64MultiArray>::SharedPtr detection_sub_;
    struct DetectionInfo
    {
        float x1;
        float y1;
        float x2;
        float y2;
        float confidence;
        int64_t stamp_ns; // 改为存储纳秒时间戳
    };

    std::vector<DetectionInfo> latest_detections_;
    std::mutex detection_mutex_;
};

bool bind_to_cpu(int cpu_core)
{
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpu_core, &cpuset);
    pid_t pid = getpid();
    if (sched_setaffinity(pid, sizeof(cpu_set_t), &cpuset) == -1)
    {
        std::cerr << "Failed to bind to CPU " << cpu_core
                  << ": " << strerror(errno) << std::endl;
        return false;
    }
    std::cout << "Process bound to CPU " << cpu_core << std::endl;
    return true;
}

int main(int argc, char **argv)
{
    int cpu_core = -1;
    for (int i = 1; i < argc; ++i)
    {
        if (strcmp(argv[i], "--cpu") == 0 && i + 1 < argc)
        {
            cpu_core = atoi(argv[i + 1]);
            for (int j = i; j + 2 <= argc; ++j)
                argv[j] = argv[j + 2];
            argc -= 2;
            break;
        }
    }
    if (cpu_core != -1 && !bind_to_cpu(cpu_core))
    {
        std::cerr << "Continuing without CPU binding" << std::endl;
    }
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<AstraNode>());
    rclcpp::shutdown();
    return 0;
}
